import { Annotation } from "@langchain/langgraph";
import { TaskSchedulerNode, WorkTask } from "./nodes/taskSchedulerNode";
import { UnknowReplyNode } from "./nodes/unknowReplyNode";
import { OutlineNode } from "./nodes/outlineNode";
import { ImageCmdNode } from "./nodes/imageCmdNode";
import { ImageExtNode } from "./nodes/imageExtNode";
import { ImageGenNode } from "./nodes/imageGenNode";
import { BaseWorkflowInputVars, WorkflowBase } from "../../core";

export class GenPosterGraph extends WorkflowBase {
    WorkflowVars = Annotation.Root({
        ...BaseWorkflowInputVars.spec,
        imageList: Annotation<any[]>,
        imageSize: Annotation<string>,
    })
    get WorkflowID() {
        return "GenPosterGraph";
    }

    protected _TaskSchedulerNode = new TaskSchedulerNode();
    // protected _UnknowReplyNode = new UnknowReplyNode();
    protected _OutlineNode = new OutlineNode();
    protected _ImageCmdNode = new ImageCmdNode();
    protected _ImageExtNode = new ImageExtNode();
    protected _ImageGenNode = new ImageGenNode();

    InitNode() {
        // this.AddNode(this._TaskSchedulerNode);
        // this.AddNode(this._UnknowReplyNode);
        this.AddNode(this._OutlineNode);
        this.AddNode(this._ImageCmdNode);
        this.AddNode(this._ImageExtNode);
        this.AddNode(this._ImageGenNode);
    }

    InitEdge() {
        this._StartNode.isGenTitle = false;
        this.AddEdge(this._StartNode,
            this._ImageExtNode,
            this._OutlineNode,
            this._ImageCmdNode,
            this._ImageGenNode,
            this._EndNode);
    }
}